import os
import requests
from bs4 import BeautifulSoup
from urllib.parse import urlparse
import trafilatura
import logging

class WebContentCrawler:
    def __init__(self, timeout=10):
        self.timeout = timeout
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }

    def extract_main_content(self, url):
        """
        使用trafilatura库提取网页主要内容
        :param url: 网页URL
        :return: 提取的主要文本内容
        """
        try:
            # 使用requests获取网页内容
            response = requests.get(url, headers=self.headers, timeout=self.timeout)
            response.raise_for_status()

            # 使用trafilatura提取主要内容
            extracted_content = trafilatura.extract(
                response.content.decode('utf-8', errors='replace'),  # 使用 UTF-8 解码
                include_comments=False,
                include_tables=False,
                no_fallback=True
            )
            print("获取搜索的内容: " + extracted_content)

            # 内容长度控制和清洗
            if extracted_content:
                # 限制内容长度，避免超长
                return extracted_content[:100000].strip()
            return "无法提取有效内容"

        except Exception as e:
            print(f"内容提取错误 {url}: {e}")
            return f"内容提取失败: {str(e)}"

    def process_search_results(self, search_results):
        """
        为搜索结果添加网页主要内容
        :param search_results: 搜索结果列表
        :return: 增强后的搜索结果
        """
        enhanced_results = []
        for result in search_results:
            try:
                # 提取网页内容
                content = self.extract_main_content(result['link'])
                
                enhanced_result = result.copy()
                enhanced_result['content'] = content
                enhanced_results.append(enhanced_result)
            
            except Exception as e:
                print(f"处理 {result['link']} 时出错: {e}")
                enhanced_results.append(result)
        
        return enhanced_results

    def enhance_ollama_search_prompt(slef, query, enhanced_results):
        """
        根据获取的网页内容生成更精确的提示词
        :param query: 原始查询
        :param enhanced_results: 包含网页内容的搜索结果
        :return: 增强型提示词
        """
        prompt = f"针对问题： '{query}', 以下是相关网页内容：\n\n"
        
        for idx, result in enumerate(enhanced_results, 1):
            prompt += f"来源 {idx}: {result['title']}\n"
            prompt += f"链接: {result['link']}\n"
            prompt += f"摘要: {result['snippet']}\n"
            prompt += f"主要内容: {result['content']}...\n\n"
        
        prompt += "请基于这些详细信息，提供一个全面、准确的回答。注意整合不同来源的信息，并保持客观中立。"
        
        return prompt

# 使用示例
def test_web_content_crawler():
    crawler = WebContentCrawler()
    
    resp = crawler.extract_main_content('https://news.un.org/zh/story/2024/01/1125512')
    
    print(resp)

if __name__ == '__main__':
    test_web_content_crawler()